#!/usr/bin/env python
import nltk
import csv
import sys
import time
import calendar
from math import log
from svmutil import *

NUM_DOCS = 3000
TOKEN_FILE = 'tokens.txt'
DATA_FILE = 'data.csv'
LABEL_TO_CLASS = {'not a real question':0,'not constructive':1,'off topic':2,'open':3,'too localized':4}
DEBUG = False

def get_tokens(filename):
    with open(filename, 'r') as token_file:
        splits = [s[:-1].split(':') for s in token_file.readlines()]
        token_to_pmi = {s[0]:s[1] for s in splits}
        tokens = [s[0] for s in splits]
        return tokens, token_to_pmi
def tokenize(data):
    for i,row in enumerate(data):
        if DEBUG:
            print 'Tokenizing %d of %d documents'%(i,len(data))
        row['tokens'] = nltk.word_tokenize(row['BodyMarkdown'].lower())
def get_data(filename):
    if DEBUG:
        print 'Getting data from %s'%filename
    with open(filename, 'r') as csvfile:
        return [row for row in csv.DictReader(csvfile, delimiter=',', quotechar='"')]
def str_to_epoch(str):
    return calendar.timegm(time.strptime(str, '%m/%d/%Y %H:%M:%S'))
def get_features(row, tokens, token_to_pmi):
    #return [int(t in row['tokens']) for t in tokens] + [str_to_epoch(row['PostCreationDate'])-str_to_epoch(row['OwnerCreationDate']),len(row['BodyMarkdown']),len(row['Title'])]
    return [str_to_epoch(row['PostCreationDate'])-str_to_epoch(row['OwnerCreationDate']),len(row['BodyMarkdown']),len(row['Title'])]

if __name__ == '__main__':
    data = get_data(sys.argv[1])[:NUM_DOCS]
    tokenize(data)
    tokens,token_to_pmi = get_tokens(TOKEN_FILE)
    all_features = []
    all_classes = {l:[] for l in LABEL_TO_CLASS}
    with open(DATA_FILE, 'w') as data_file:
        for i,row in enumerate(data):
            if DEBUG:
                print 'Getting features for %d of %d documents'%(i,len(data))
            all_features.append(get_features(row,tokens,token_to_pmi))
            for key,val in all_classes.items():
                all_classes[key].append(1 if row['OpenStatus'] == key else 0)

    for key,val in all_classes.items():
        m = svm_train(val, all_features, '-q -b 1') 
        p_label, p_acc, p_val = svm_predict(val, all_features, m)
        print p_val
